
import numpy as np
import matplotlib.pyplot as plt
import os

Ts,Pstrat,h2ocol,fv,whichrun,trun,gastocld,cldtogas,tohaze,hazeprod,hazeloss,hazeconc,zstrat,prodmoist,h2o_top,prodstrat,moistbot = [],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[]
mypath = 'outputs_1e5_jan6_shield300_rain_10ppm_format/'

for ifile in os.listdir(mypath):

    if '-fv-' in ifile:
        continue
    if not ifile.endswith('out-1'):
        continue
    
    with open(mypath+ifile,'r') as file:
        lines = file.readlines()
    okay = 0

    Ts += [ifile.split('Ts')[1][:-7]]
    fv += [ifile.split('fv')[1].split('Ts')[0][:-1]]
    
    # make sure run converged
    timelines = [i for i in lines if 'ELAPSED TIME' in i]
    try:
        tottime = timelines[-1].split()[3]
    except:
        tottime = 0.0
        
    if float(tottime)>1e7:
        print(ifile,fv[-1],Ts[-1], "converged")
        a=0
    else:
        print("this run did not converge ",ifile,fv[-1],Ts[-1])
        prodstrat += [-1]
        prodmoist += [-1]
        moistbot += [0]
        Pstrat += [0]
        h2ocol +=[0]
        h2o_top += [0]
        zstrat += [0]
        trun += [0]
        continue
        
    trun += [float(tottime)]
    
    # search for Pstrat
    indices = [index for index, line in enumerate(lines) if 'EDDY' in line][0]
    z,p,d = [],[],[]
    for iz in range(indices+3,indices+113):
        line = lines[iz]
        if line.split()[2] == '200.0':
            Pstrat += [line.split()[4]]
            zstrat += [float(line.split()[1])]
            break
    if len(Pstrat) < len(Ts):
        Pstrat += [float('NaN')]
        zstrat += [1000.0]

    for iz in range(indices+3,indices+113):
        line = lines[iz]
        z += [float(line.split()[1])]
        p += [float(line.split()[4])]
        d += [float(line.split()[3])]

    # search for lower layer moist 
    check = 'V          HUM'
    tempbot=0
    indices = [index for index,line in enumerate(lines) if check in line][0]
    for iz in range(indices+1,indices+111):
        line = lines[iz]
        if (line.split()[4] == '1.000E+00') or (float(line.split()[4]) > 0.9999):
            temp = float(line.split()[1])
            indx = np.where(np.array(z) == temp)
            tempbot = float(line.split()[1])

    #for txt uncomment two lines 
    if tempbot >0:
        moistbot += [tempbot]
    if len(moistbot) < len(Ts):
        moistbot += [0.0]
        
    check = 'ALTITUDE   H2O        N2         CO2 '
    indices = [index for index,line in enumerate(lines) if check in line][0]
    h2o = []
    ctr = 0
    for iz in range(indices+1,indices+111):
        line = lines[iz].split()
        h2o += [float(line[2])/d[ctr]]
        ctr += 1
    line = lines[iz+2].split()
    h2ocol += [float(line[1])]
    h2o_top += [h2o[-1]]

    # search for loss of SO2 and H2SO4 to cloud
    check = 'RATE( 311) RATE( 312)'
    indx = [index for index,line in enumerate(lines) if check in line][0]
    line = lines[indx+112]
    stemp,mtemp = [],[]
    for iz in range(indx+1,indx+110):
        line = lines[iz].split()
        dz = (float(z[iz-indx])-float(z[iz-indx-1]))*1e5 # cm to km
        if moistbot[-1] < 1000.0:
            if (float(line[1]) <= moistbot[-1]):
                mtemp += [float(line[2])*dz]
            elif (float(line[1]) > moistbot[-1]):
                stemp += [float(line[2])*dz]
        else:
            stemp += [float(line[2])*dz]
            mtemp += [0.0]
    prodstrat += [sum(stemp)]
    prodmoist += [sum(mtemp)]

    
Ts = [float(i) for i in Ts]
fv = [float(i) for i in fv]
Ts = np.array(Ts)
fv = np.array(fv)
prodmoist = np.array(prodmoist)
trun = np.array(trun)
prodstrat = np.array(prodstrat)
h2o_top = np.array(h2o_top)
Ts_unique = np.unique(Ts)
fv_unique = np.unique(fv)
Ts_sorted = np.sort(Ts_unique)
fv_sorted = np.sort(fv_unique)
h2ocol = np.array(h2ocol)
moistbot = np.array(moistbot)


# sort haze conc
new_strat = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_moist = np.zeros((len(Ts_sorted),len(fv_sorted)))
newheight = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_time = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_Ts = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_fv = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_h2o = np.zeros((len(Ts_sorted),len(fv_sorted)))
new_h2o_col = np.zeros((len(Ts_sorted),len(fv_sorted)))

for i in range(0,len(Ts_sorted)):
    for j in range(0,len(fv_sorted)):
        indx = np.where( (Ts == Ts_sorted[i]) & (fv == fv_sorted[j]) )
        print(i,j,indx[0],len(indx[0]))
        if len(indx[0])==1:
            if trun[indx] > 1e7:
                new_time[i,j] = 1.0 # trun[indx]
            else:
                new_time[i,j] = 0.0
            new_strat[i,j] = prodstrat[indx]
            new_moist[i,j] = prodmoist[indx]
            newheight[i,j] = moistbot[indx]
            new_Ts[i,j] = Ts[indx]
            new_fv[i,j] = fv[indx]
            new_h2o[i,j] = h2o_top[indx]
            new_h2o_col[i,j] = h2ocol[indx]
        else:
            continue

plt.clf()
plt.imshow(np.log10(new_strat))
plt.colorbar()
plt.savefig('strat.png')

plt.clf()
plt.imshow(new_time)
plt.colorbar()
plt.savefig('time.png')

plt.clf()
plt.imshow(np.log10(new_moist))
plt.colorbar()
plt.savefig('moist.png')

plt.clf()
plt.imshow(np.log10(new_h2o))
cbar = plt.colorbar()
x_labels = ['1e-6','', '1e-4','', '5e-4','', '2e-3','', '1e-2','', '6e-2','', '3e-1','']
y_labels = ['300','','400','','500','','600','','700','','800','','900','','1000']
plt.xticks(ticks=np.arange(len(x_labels)), labels=x_labels)
plt.yticks(ticks=np.arange(len(y_labels)), labels=y_labels)
plt.xlabel('fv')
plt.ylabel('Ts')
cbar.set_label('log10 Xh2o')

plt.savefig('h2otop.png')

plt.clf()
b = 0.202 # log10(D12 x=1/2) = -0.6947 from https://srd.nist.gov/jpcrdreprint/1.3253094.pdf
b = 2e18 # matlab
# from wordsworth 2013
f = new_h2o
mu = ((new_h2o*18.0)+((1-(new_h2o))*44.0))#/2.0
g = 887.0
kb = 1.38e-16
T = 260.0
Hn = kb*T/(mu*g*6.67e-24)
Hw = kb*T/(18.0*g*6.67e-24)
flux = b*f*((1/Hn)-(1/Hw))
print('b f Hn Hw mu flux')
print(b,f[10,10],Hn[10,10],Hw,mu[10,10],flux[10,10])
plt.imshow(np.log10(flux)) # 1e27/flux
plt.colorbar()
plt.xticks(ticks=np.arange(len(x_labels)), labels=x_labels)
plt.yticks(ticks=np.arange(len(y_labels)), labels=y_labels)
plt.xlabel('fv')
plt.ylabel('Ts')
plt.savefig('tauloss.png')

plt.clf()
plt.imshow(np.log10(new_h2o_col/(flux*3e7)))
tau = new_h2o_col/(flux*3e7)
import pickle as pickle
with open('tau.pickle','wb') as handle:
    pickle.dump(tau,handle)#x,protocol=pickle.HIGHEST_PRIORITY)
plt.colorbar()
plt.xticks(ticks=np.arange(len(x_labels)), labels=x_labels)
plt.yticks(ticks=np.arange(len(y_labels)), labels=y_labels)
plt.xlabel('fv')
plt.ylabel('Ts')
plt.savefig('tauh2o.png')

print(new_strat)

plt.clf()
plt.imshow(np.log10(newheight))
plt.colorbar()
plt.savefig('heightmoist.png')
#xexit()


# write to text file for simplicity
outfile = open('summary_kineticsresults_eddy1e5_shield300_krON_jan25.txt','w')
for i in range(0,len(new_Ts.flatten())):
    temp = '%.3e'%new_Ts.flatten()[i]+'  '+'%.3e'%new_fv.flatten()[i]+'  '+'%.3e'%new_strat.flatten()[i]+'  '+'%.3e'%new_moist.flatten()[i]+'\n'
    outfile.write(temp)
outfile.close()
exit()
#Ts,Pstrat,fv,gastocld,cldtogas,tohaze,hazeprod,hazeloss,hazeconc

# prepare to plot
Ts = np.array(Ts)
fv = np.array(fv)
hazeconc = np.array(hazeconc)

Ts_unique = np.unique(Ts)
fv_unique = np.unique(fv)

Ts_grid, fv_grid = np.meshgrid(Ts_unique, fv_unique)
hazeconc_grid = np.reshape(hazeconc, Ts_grid.shape)
print(np.shape(Ts_grid),np.shape(fv_grid))
print(np.shape(hazeconc_grid))
plt.imshow(np.log10(hazeconc_grid))
plt.colorbar()
#contour = plt.contourf(Ts_grid, fv_grid, hazeconc_grid, cmap='viridis')
plt.savefig('hazeconc.png')

'''plt.contourf(Ts_a,np.log10(f_v_a),p_strat_a/1e5,levels = [0,0.01,0.05,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8])
plt.xlabel('surface temperature [K]')
plt.ylabel('log10 of surface H2O molar conc. [mol mol-1]')
plt.title('Tropopause pressure [bar]')
plt.colorbar()
'''
